#include "CudaHelpers.h"
#include "DataLoader.h"
#include <cuda.h>
#include <CImg.h>
#include "cublas.h"

using namespace cimg_library;

#define MIN(a, b)  (((a) < (b)) ? (a) : (b))
#define MAX(a, b)  (((a) > (b)) ? (a) : (b))

float* AllocateAndCopy(float* h, int size)
{
	float* d = 0;
	cudaMalloc((void**)&d, size * sizeof(float));
	cublasStatus status = cublasSetVector(size, sizeof(float), h, 1, d, 1);
    if (status != CUBLAS_STATUS_SUCCESS) {
        fprintf (stderr, "!!!! device access error (write)\n");
        getchar();
		exit(-1);
    }

	return d;
}

float* Allocate(int size)
{
	float* d = 0;
	cublasStatus status = cublasAlloc(size, sizeof(float), (void**)&d);
    if (status != CUBLAS_STATUS_SUCCESS) {
        fprintf (stderr, "!!!! device memory allocation error\n");
        getchar();
		exit(-1);
    }
	cudaMemset(d, 0, size * sizeof(float));
	return d;
}

void Free(float* d)
{
	cudaError_t status = cudaFree(d);
    if (status != cudaSuccess) {
		fprintf (stderr, "!!!! memory free error: %s (%d)\n", cudaGetErrorString(status), status);
        getchar();
		exit(-1);
    }
}

void SaveImageData(string filename, float* data, int width, int height)
{
	unsigned char* d = (unsigned char*)malloc(width * height * 4 * sizeof(unsigned char));
	for(int i = 0;i < width * height * 4;i++)
	{
		d[i] = (unsigned char)(255.f - MAX(.0f, MIN(1.0f, data[i])) * 255.f);
		//printf("%d ", d[i]);
	}

	CImg<unsigned char> image(d, width, height, 1, 3);

	image.save(filename.c_str());

	free(d);
}

void SaveAsImage(string filename, float* d, int width, int height, float offset, float scale)
{
	float* h = (float*)malloc(width * height * 4 * sizeof(float));
	cublasStatus status = cublasGetVector(width * height * 4, sizeof(float), d, 1, h, 1);
	if (status != CUBLAS_STATUS_SUCCESS) {
        fprintf (stderr, "!!!! get vector error (%d)\n", status);
        getchar();
		exit(-1);
    }

	for(int i = 0;i < width;i++)
		for(int j = 0;j < height;j++)
			for(int k = 0;k < 4;k++) {
				//printf("%f ", h[i * height * 4 + j * 4 + k]);
				h[k * width * height + j * width + i] = h[k * width * height + j * width + i] * scale + offset;
			}

	SaveImageData(filename, h, width, height);
	free(h);
}
